package com.shuai.AI.demo.rag;

import java.util.List;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

/** 查询扩展器 Demo */
@Component
public class MultiQueryExpanderDemo {

  private final ChatClient.Builder chatClientBuilder;

  public MultiQueryExpanderDemo(ChatModel dashscopeChatModel) {
    this.chatClientBuilder = ChatClient.builder(dashscopeChatModel);
  }

  public List<Query> expand(String query) {
    MultiQueryExpander queryExpander =
        MultiQueryExpander.builder()
            .chatClientBuilder(chatClientBuilder)
            .numberOfQueries(3)
            .build();
    List<Query> queries = queryExpander.expand(new Query(query));
    return queries;
  }
}
